package show.lmm.spring_data_r2dbc.core;

import io.r2dbc.spi.ConnectionFactories;
import io.r2dbc.spi.ConnectionFactory;
import org.springframework.r2dbc.connection.lookup.AbstractRoutingConnectionFactory;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

/**
 * 动态数据源
 */
public class DynamicDataSource extends AbstractRoutingConnectionFactory {

    // 数据源列表
    private final Map<String, ConnectionFactory> connectionFactoryMap = new HashMap<>();

    // 数据库名称key
    public static final String CURRENT_DB_NAME = "dbName";

    /**
     * 查询首个db链接
     */
    public ConnectionFactory getFirstConnection() {
        if (connectionFactoryMap.isEmpty()) {
            return null;
        }
        final String key = new ArrayList<>(connectionFactoryMap.keySet()).get(0);
        return connectionFactoryMap.get(key);
    }

    /**
     * 设置可用数据源
     */
    public void setTargetConnectionFactories() {
        super.setTargetConnectionFactories(connectionFactoryMap);
    }

    @Override
    public void setDefaultTargetConnectionFactory(Object defaultTargetConnectionFactory) {
        super.setDefaultTargetConnectionFactory(defaultTargetConnectionFactory);
    }

    @Override
    protected Mono<Object> determineCurrentLookupKey() {
        return Mono.deferContextual(Mono::just).handle(((context, sink) -> {
            final String dbName = context.getOrDefault(CURRENT_DB_NAME, null);
            if (Optional.ofNullable(dbName).isPresent()) {
                sink.next(dbName);
            }
        }));
    }

    /**
     * 新增数据源
     *
     * @param dbName        数据库名称
     * @param connectionUrl 链接字符串
     */
    public void addDataSource(String dbName, String connectionUrl) {
        connectionFactoryMap.put(dbName, ConnectionFactories.get(connectionUrl));
    }
}